// LoadNTDriver.cpp: implementation of the CLoadNTDriver class.
//
//////////////////////////////////////////////////////////////////////

#include "stdafx.h"
#include "LoadNTDriver.h"

#ifdef _DEBUG
#undef THIS_FILE
static char THIS_FILE[] = __FILE__;
#define new DEBUG_NEW
#endif

#define				SECURITY_STRING_LEN							168
#define				LG_PAGE_SIZE								4096
#define				MAX_KEY_LENGTH								1024
#define				LG_SLEEP_TIME								4000

const BYTE g_szSecurity[SECURITY_STRING_LEN] =
{
	0x01,0x00,0x14,0x80,0x90,0x00,0x00,0x00,0x9c,0x00,0x00,0x00,0x14,0x00,0x00,0x00,0x30,0x00,0x00,0x00,0x02,
	0x00,0x1c,0x00,0x01,0x00,0x00,0x00,0x02,0x80,0x14,0x00,0xff,0x01,0x0f,0x00,0x01,0x01,0x00,0x00,0x00,0x00,
	0x00,0x01,0x00,0x00,0x00,0x00,0x02,0x00,0x60,0x00,0x04,0x00,0x00,0x00,0x00,0x00,0x14,0x00,0xfd,0x01,0x02,
	0x00,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x05,0x12,0x00,0x00,0x00,0x00,0x00,0x18,0x00,0xff,0x01,0x0f,0x00,
	0x01,0x02,0x00,0x00,0x00,0x00,0x00,0x05,0x20,0x00,0x00,0x00,0x20,0x02,0x00,0x00,0x00,0x00,0x14,0x00,0x8d,
	0x01,0x02,0x00,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x05,0x0b,0x00,0x00,0x00,0x00,0x00,0x18,0x00,0xfd,0x01,
	0x02,0x00,0x01,0x02,0x00,0x00,0x00,0x00,0x00,0x05,0x20,0x00,0x00,0x00,0x23,0x02,0x00,0x00,0x01,0x01,0x00,
	0x00,0x00,0x00,0x00,0x05,0x12,0x00,0x00,0x00,0x01,0x01,0x00,0x00,0x00,0x00,0x00,0x05,0x12,0x00,0x00,0x00
};

//////////////////////////////////////////////////////////////////////
// Construction/Destruction
//////////////////////////////////////////////////////////////////////

CLoadNTDriver::CLoadNTDriver(const char* szDriverName,const char* szDriverFullPath)
{
	m_szDriverName = szDriverName;
	m_szDriverFullPath = szDriverFullPath;
}

CLoadNTDriver::~CLoadNTDriver()
{

}

bool CLoadNTDriver::InstallDriver()
{
	char	szBuf[LG_PAGE_SIZE] = {0};
	HKEY	hKey;
	DWORD	dwData = 0;

	if( NULL==m_szDriverName || NULL==m_szDriverFullPath )
	{
		return false;
	}
	memset(szBuf,0,LG_PAGE_SIZE);
	strcpy_s(szBuf, "SYSTEM\\CurrentControlSet\\Services\\");
	strcat_s(szBuf,m_szDriverName);
	if(RegCreateKeyEx(HKEY_LOCAL_MACHINE,szBuf,0,"",0,KEY_ALL_ACCESS,NULL,&hKey,(LPDWORD)&dwData)!=ERROR_SUCCESS)
	{
		return false;
	}
	strcpy_s(szBuf,m_szDriverName);
	if(RegSetValueEx(hKey,"DisplayName",0,REG_SZ,(CONST BYTE*)szBuf,(DWORD)strlen(szBuf))!=ERROR_SUCCESS)
	{
		return false;
	}
	dwData = 1;
	if(RegSetValueEx(hKey,"ErrorControl",0,REG_DWORD,(CONST BYTE*)&dwData,sizeof(DWORD))!=ERROR_SUCCESS)
	{
		return false;
	}
	strcpy_s(szBuf,"\\??\\");
	strcat_s(szBuf,m_szDriverFullPath);
	if(RegSetValueEx(hKey,"ImagePath",0,REG_SZ,(CONST BYTE*)szBuf,(DWORD)strlen(szBuf))!=ERROR_SUCCESS)
	{
		return false;
	}
	dwData = 3;
	if(RegSetValueEx(hKey,"Start",0,REG_DWORD,(CONST BYTE*)&dwData,sizeof(DWORD))!=ERROR_SUCCESS)
	{
		return false;
	}
	dwData = 1;
	if(RegSetValueEx(hKey,"Type",0,REG_DWORD,(CONST BYTE*)&dwData,sizeof(DWORD))!=ERROR_SUCCESS)
	{
		return false;
	}
	RegFlushKey(hKey);
	RegCloseKey(hKey);
	strcpy_s(szBuf,"SYSTEM\\CurrentControlSet\\Services\\");
	strcat_s(szBuf,m_szDriverName);
	strcat_s(szBuf,"\\Security");
	if(RegCreateKeyEx(HKEY_LOCAL_MACHINE,szBuf,0,"",0,KEY_ALL_ACCESS,NULL,&hKey,(LPDWORD)&dwData)!=ERROR_SUCCESS)
	{
		return false;
	}
	dwData = SECURITY_STRING_LEN;
	if(RegSetValueEx(hKey,"Security",0,REG_BINARY,g_szSecurity,dwData)!=ERROR_SUCCESS)
	{
		return false;
	}
	RegFlushKey(hKey);
	RegCloseKey(hKey);

	return true;
}

bool CLoadNTDriver::StartDriver()
{
	SC_HANDLE		schManager;
	SC_HANDLE		schService;
	SERVICE_STATUS	svcStatus;
	bool			bStarted=false;
	int				i = 0;

	if(NULL == m_szDriverName)
	{
		return false;
	}
	if(!CreateDriver())
	{
		return false;
	}
	schManager = OpenSCManager(NULL,NULL,SC_MANAGER_ALL_ACCESS);
	if(NULL == schManager)
	{
		return false;
	}
	schService = OpenService(schManager,m_szDriverName,SERVICE_ALL_ACCESS);
	if(NULL == schService)
	{
		CloseServiceHandle(schManager);
		return false;
	}
	if(ControlService(schService,SERVICE_CONTROL_INTERROGATE,&svcStatus))
	{
		if(svcStatus.dwCurrentState == SERVICE_RUNNING)
		{
			CloseServiceHandle(schService);
			CloseServiceHandle(schManager);
			return true;
		}
	}
	else if(GetLastError()!=ERROR_SERVICE_NOT_ACTIVE)
	{
		CloseServiceHandle(schService);
		CloseServiceHandle(schManager);
		return false;
	}
	if(!StartService(schService,0,NULL))
	{
		CloseServiceHandle(schService);
		CloseServiceHandle(schManager);
		return false;
	}
	for(i=0; i<10; i++)
	{
		if( ControlService(schService,SERVICE_CONTROL_INTERROGATE,&svcStatus) && svcStatus.dwCurrentState==SERVICE_RUNNING )
		{
			bStarted = true;
			break;
		}
		Sleep(LG_SLEEP_TIME);
	}
	CloseServiceHandle(schService);
	CloseServiceHandle(schManager);

	return bStarted?true:false;
}

bool CLoadNTDriver::CreateDriver()
{
	SC_HANDLE		schManager;
	SC_HANDLE		schService;
	SERVICE_STATUS	svcStatus;
	bool			bStopped = false;
	int				i = 0;

	if( NULL == m_szDriverName || NULL == m_szDriverFullPath )
	{
		return false;
	}
	schManager = OpenSCManager(NULL,NULL,SC_MANAGER_ALL_ACCESS);
	if(NULL == schManager)
	{
		return false;
	}
	schService = OpenService(schManager,m_szDriverName,SERVICE_ALL_ACCESS);
	if(NULL != schService)
	{
		if(ControlService(schService,SERVICE_CONTROL_INTERROGATE,&svcStatus))
		{
			if(svcStatus.dwCurrentState!=SERVICE_STOPPED)
			{
				if(0 == ControlService(schService,SERVICE_CONTROL_STOP,&svcStatus))
				{
					CloseServiceHandle(schService);
					CloseServiceHandle(schManager);
					return false;
				}
				for(i=0; i<10; i++)
				{
					if( ControlService(schService,SERVICE_CONTROL_INTERROGATE,&svcStatus)==0 || svcStatus.dwCurrentState==SERVICE_STOPPED )
					{
						bStopped = true;
						break;
					}
					Sleep(LG_SLEEP_TIME);
				}
				if(!bStopped)
				{
					CloseServiceHandle(schService);
					CloseServiceHandle(schManager);
					return false;
				}
			}
		}
		CloseServiceHandle(schService);
		CloseServiceHandle(schManager);
		return true;
	}
	schService = CreateService(schManager,m_szDriverName,m_szDriverName,SERVICE_ALL_ACCESS,SERVICE_KERNEL_DRIVER,SERVICE_SYSTEM_START,SERVICE_ERROR_NORMAL,m_szDriverFullPath,NULL,NULL,NULL,NULL,NULL);
	if(NULL == schService)
	{
		CloseServiceHandle(schManager);
		return false;
	}
	CloseServiceHandle(schService);
	CloseServiceHandle(schManager);

	return true;
}

bool CLoadNTDriver::StopDriver()
{
	SC_HANDLE		schManager;
	SC_HANDLE		schService;
	SERVICE_STATUS	svcStatus;
	bool			bStopped = false;
	int				i = 0;

	schManager = OpenSCManager(NULL,0,0);
	if(NULL == schManager)
	{
		return false;
	}
	schService = OpenService(schManager,m_szDriverName,SERVICE_ALL_ACCESS);
	if(NULL == schService)
	{
		CloseServiceHandle(schManager);
		return false;
	}
	if(ControlService(schService,SERVICE_CONTROL_INTERROGATE,&svcStatus))
	{
		if(svcStatus.dwCurrentState!=SERVICE_STOPPED)
		{
			if(!ControlService(schService,SERVICE_CONTROL_STOP,&svcStatus))
			{
				CloseServiceHandle(schService);
				CloseServiceHandle(schManager);
				return false;
			}
			for(i=0; i<10; i++)
			{
				if( ControlService(schService,SERVICE_CONTROL_INTERROGATE,&svcStatus)==0 || svcStatus.dwCurrentState==SERVICE_STOPPED )
				{
					bStopped = true;
					break;
				}
				Sleep(LG_SLEEP_TIME);
			}
			if(!bStopped)
			{
				CloseServiceHandle(schService);
				CloseServiceHandle(schManager);
				return false;
			}
		}
	}
	CloseServiceHandle(schService);
	CloseServiceHandle(schManager);

	return true;
}

bool CLoadNTDriver::RemoveDriver()
{
	HKEY hKey;
	char szBuf[LG_PAGE_SIZE] = {0};
	char szDriverName[MAX_PATH] = {0};

	memset(szBuf,0,LG_PAGE_SIZE);
	memset(szDriverName,0,MAX_PATH);
	strcpy_s(szDriverName,m_szDriverName);
	strcpy_s(szBuf,"SYSTEM\\CurrentControlSet\\Services\\");
	if(RegOpenKeyEx(HKEY_LOCAL_MACHINE,szBuf,0,KEY_ALL_ACCESS,&hKey)!=ERROR_SUCCESS)
	{
		return false;
	}
	if((RegDeleteKeyNT(hKey,szDriverName))!=ERROR_SUCCESS)
	{
		return false;
	}
	RegCloseKey(hKey);

	return 0;
}

long CLoadNTDriver::RegDeleteKeyNT(HKEY hStartKey,LPTSTR pKeyName)
{
	DWORD	dwSubKeyLength = 0;
	LPTSTR	pSubKey = NULL;
	TCHAR	szSubKey[MAX_KEY_LENGTH] = {0};
	HKEY	hKey;
	LONG	lRet = 0;

	if( pKeyName && lstrlen(pKeyName) )
	{
		if((lRet = RegOpenKeyEx(hStartKey,pKeyName,0,KEY_ENUMERATE_SUB_KEYS|DELETE,&hKey))==ERROR_SUCCESS)
		{
			while(lRet == ERROR_SUCCESS)
			{
				dwSubKeyLength = MAX_KEY_LENGTH;
				lRet = RegEnumKeyEx(hKey,0,szSubKey,(LPDWORD)&dwSubKeyLength,NULL,NULL,NULL,NULL);
				if(lRet == ERROR_NO_MORE_ITEMS)
				{
					lRet = RegDeleteKey(hStartKey,pKeyName);
					break;
				}
				else if(lRet == ERROR_SUCCESS)
				{
					lRet = RegDeleteKeyNT(hKey,szSubKey);
				}
			}
			RegCloseKey(hKey);
		}
	}
	else
	{
		lRet = ERROR_BADKEY;
	}

	return lRet;
}